#
#   SolveHarmonicOscillator.py
#

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from FunctionalBohm import SchrodingerNDT
from math import pi

import os
if not os.path.isdir("./data"):
    os.makedirs("./data")

def segments(min,max,steps):
    return [i/float(steps)*(max-min)+min for i in range(steps+1)]

# Start solving:

from HarmonicOscillator1D import Psi as wvfn
from HarmonicOscillator1D import V
from HarmonicOscillator1D import mass
from HarmonicOscillator1D import hbar
from HarmonicOscillator1D import omega

Schrod = SchrodingerNDT(V,wvfn,mass,dimensions=1,dx=0.000005) # we need to specify a small enough stepsize!!!

# Alternatively we could have defined initial conditions of interest while defining Schrod, eg
# Schrod = SchrodingerNDT(..., positions=segments(-2,2,100), time=1.3, ...)
# But that's kind of boring because the positions are fixed, do not follow the flow (until I code that up later).

# get the physical variables of interest as functions of (t,x)

J = Schrod.get_Probability_Flux_Density()
H = Schrod.get_Hamiltonian()
Q = Schrod.get_Quantum_Potential()
#traj = S.get_Trajectory(X_ic,T_ic,T_final,timesteps) # trajectory as a list.
#time = S.get_Times(T_ic,T_final,timesteps)
#p = map(S.get_momentum(),time,traj)

# We can now get Hamiltonian as a function of position (along rows) and time (along columns)

count = 0
for time in segments(0,2.0*pi/omega,100): # A full `physical' period.
    print(count)
    outfile = open("./data/time" + "%06d" % (count) + ".csv",'w')
    count += 1
    for i in segments(-40,40,500):
        #outfile.write(str(i)+", "+str(Q(time,i))+", "+str(V(i))+"\n")
        #outfile.write(str(i)+", "+str(J(time,i)[0])+", "+str(V(i))+"\n")
        outfile.write(str(i)+", "+str(H(time,i))+", "+str(V(i))+", "+str(abs(wvfn(time,i))**2)+"\n")
        
    outfile.close()

print("Updated './data/'")

## To find level sets at height h, solve H - h == 0:
#
#from scipy.optimize import fsolve
#
#def Levelset(H,t,h,start):
#   def diff(x): return H(t,x)-h
#   return fsolve(diff,start) # we need a vector of starting guesses...
#
#print(Levelset(H,t,6.25,[-0.75,-0.15,0.1,0.6]))
#
## NB: It would be nice to find the level sets without the starting guesses...
## Maybe just use segments(-2,2,400) as the guesses and then clean up duplicates later?
